[lora] feat: multi lora async (full PR, see stacked diff PRs instead)#1638
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fully-async multi-LoRA training framework, enabling concurrent training of multiple LoRA adapters against a shared base model. It adds a background rollout worker, an async round-robin data source, a Ray-based controller with an HTTP control-plane API, and integrates slot-keyed adapter management into the Megatron-LM and SGLang backends. The review feedback identifies a critical distributed deadlock risk in actor.py caused by non-deterministic set iteration during collective operations, which can be fixed by sorting. Additionally, the reviewer recommends replacing several assert statements with ValueError for input validation and closing the httpx.Client in the smoke test to prevent resource leaks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
yushengsu-thu
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
Ran a deep verification pass over this branch together with its two sglang dependency PRs (sgl-project#30912/#30913); every finding below was adversarially re-verified against the code. Two issues to flag on the miles side — one is a hard crash/correctness bug in the logprob-recompute path, the other is a lifecycle race that silently corrupts rollout data across adapter slot reuse. (Side note: the deeper sglang-side findings are filed on the two dependency PRs directly.)
_build_prefill_scoring_payload keyed LoRA solely on is_lora_enabled(args)
(always true in multi-LoRA) and sent the single-adapter name miles_lora,
which is never registered on multi-LoRA engines -> every step crashed with
'adapters are not loaded' when --recompute-logprobs-via-prefill is on; the
batched path also applied payloads[0]'s lora_path to a mixed-adapter batch.
- resolve lora_path per sample: adapter samples score under their own
__miles_slot_{N}, single-adapter LoRA keeps miles_lora, base keeps none
- batched scoring groups by (logprob_start_len, lora_path) and rejects
mixed-adapter batch payloads
- centralize the engine-side slot adapter name as slot_lora_name(), shared
by rollout, prefill scoring, and the weight-push mixin
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
The retire-time prefix abort fires exactly once (RETIRING->CLEANUP); requests
can survive it (a multi-turn group POSTs its next turn after the round, or a
request sits in the engine's tokenizer window) and, once the slot is reused
and the next adapter's weights are upserted into the same __miles_slot_{N},
keep decoding under the wrong weights with no error. Batch-time collection
filters already drop such groups from training data; these changes shrink the
window and stop wasting decode on them:
- MultiLoRABackend.free_slot: re-run the prefix abort right before the slot
is released for reuse (second round, after a full step of settling)
- sglang_rollout.generate: refuse to POST for an adapter that is no longer
sampleable (deregistered/cleaned up) and abort the sample locally instead
- the tokenizer-window escape itself is fixed engine-side (sgl-project#30912)
Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
|
@mathewjhan |
Re-queuing aborted groups (add_samples) is intended but deferred to a future PR (confirmed by mathewjhan; left unimplemented for lack of testing time). strict=True so wiring it up forces marker removal. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
|
Thanks for confirming — marked |
Each sglang tokenizer worker process holds its own LoRA registry with no cross-worker sync (sgl-project/sglang#31084), so per-step adapter upserts resolve on whichever worker the router picks and fail non-deterministically. sglang rejects the upsert at runtime anyway; failing at launch avoids burning GPU time until the first weight push. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
fdefa55 to
f115610
Compare
…_step The 28-line per-slot step + per-adapter scheduler + adapter_lr logging arm moves to multi_lora_utils.step_stepped_adapter_slots; train_one_step keeps a one-call branch. The duplicated zero-grad block (prologue + release) becomes a local _zero_grads helper. No behavior change.
…s helpers train_actor's mark-batch-trained block, save_model's due-adapter checkpoint selection, and update_weights' stale-push set / version-bump bookkeeping move to multi_lora_utils (commit_trained_batch, save_due_adapter_checkpoints, select_adapters_to_push, commit_weight_push). Each call site keeps a short guarded call; comments move with the bodies. No behavior change.
…single-LoRA transport Restore _update_lora_weight_implementation to its main form (no new parameters, no behavior change for single-LoRA) and give multi-LoRA its own _update_multi_lora_weight_implementation carrying the per-slot name, the adapter's own config, the upsert RPC, and the contiguous materialization for sliced tensors. broadcast.py becomes a pure addition relative to main.
…dapter The multi-LoRA push always upserts (sglang's upsert is insert-or-overwrite, covering both first load and refresh), so the parameter threaded through _send_one_multi_lora_adapter and the implementation was never varied. Hardcode it at the RPC, gather the lazy imports at the top, and fold the adapter's r/alpha into the config with a dict merge.
The only consumer of adapter_slots is get_batch's own per-adapter token-count computation; the model.py callers declared it in their key lists but never read it. Fetch it inside get_batch (None for non-multi-LoRA runs, same as the existing dynamic_global_batch_size self-fetch) and drop it from both callers' key lists.
# Conflicts: # miles/backends/megatron_utils/bridge_lora_helpers.py
Follow the run_glm5_2_744b_a40b_lora.py convention: one run_multi_lora.py with prepare / train / full-train / serve commands over command_utils (execute_train with train_script=train_multi_lora_async.py, model registry qwen3-4B). provision.sh / run_job.sh / run_service.sh are gone; dataset and model paths move to the shared /root/datasets and /root/models layout, and the adapter yamls, smoke-client usage, and README follow.
yushengsu-thu
left a comment
yushengsu-thu
left a comment
There was a problem hiding this comment.
Summarize the comments into one or two lines. If a comment doesn't make sense, feel free to delete or remove it.
…lines Per review: every flagged comment block across the multi-LoRA files is reduced to at most two lines keeping only the load-bearing constraint; redundant restatements are deleted. One error message shortened alongside. No code change (AST-verified identical outside docstrings).
# Conflicts: # miles/backends/megatron_utils/arguments.py # miles/backends/sglang_utils/sglang_engine.py # miles/rollout/generate_utils/sample_utils.py # miles/utils/types.py
Splitting this PR into 7 stacked PRsTo reduce review load, this PR is being split into 7 module-level PRs, each an independently testable unit. This PR stays open as the integration reference (its E2E validation — 2×H200 smoke, full lifecycle — covers the whole stack). The 7 branches are pushed to Every prefix of the stack keeps
Merge order = import dependency order: everything imports 1; the trainer (3) and rollout (5, 6) talk to the controller (2); weight sync (4) uses trainer exports; the driver (7) wires it all. Verification done on each branch: import-graph completeness per layer (no branch references a module that lands later), and each layer's fast suites pass on its own tree (layer 1: 2456 passed; layer 2: 2488; layers 3–4: 266 each incl. all existing megatron_utils tests; layer 5: 400; layer 6: full Mechanics note: since these merge with squash, after each merge GitHub will still show the previous layers' changes in the next PR's diff until it is rebased — after merging k/7, rebase branches k+1..7 onto the new |
mathewjhan
changed the title
yushengsu-thu
changed the title
2 participants
Summary
Allow disagg fully async training on multiple loras (all-linear, excluding expert) per step in Miles using megatron-bridge.
Two LoRAs trained together (DAPO + GSM8K)
Feature
in_placepause generation and LoRA upsert updatesDependent PRs
SGLang
ridprefix abort: feat(sglang-miles): Support aborting requests by rid prefix - multi-lora needs sgl-project/sglang#30912Unified above: sgl-project/sglang#31253 (for testing - do not need merge)
Megatron-Bridge (radixark)
Dev setup
Tested image:
radixark/miles:dev-202607090055You need to install the 3 forks:
SGLang
This fork has two small changes to SGLang (has both PRs):
Megatron-Bridge
This fork adds a multi-lora support to Megatron-Bridge along with some helper methods/user api
Miles
Running
see:
examples/multi_loraFor normal training (not as a service):
examples/multi_lora/provision.shexamples/multi_lora/run_job.shexamples/multi_lora/run_job.sh |& tee run.logFor multi-lora training as a long running service:
examples/multi_lora/provision.shexamples/multi_lora/run_service.shexamples/multi_lora/run_service.sh |& tee run.login one shellMinor features added to existing code (backwards compatible)
AdapterRefandRewardSpectoSampletype so individual samples can access their own reward functions and adapter names during rolloutMissing features for future